"""
模型测试和预测脚本
功能：加载训练好的模型，对测试数据进行预测，生成涨幅排名结果
"""

import pandas as pd
import numpy as np
import pickle
import torch
from config import *

# 预测序列长度（使用最近32天的数据进行预测）
pred_len = 32


def inputdata(path):
    """
    读取CSV数据文件

    Args:
        path (str): CSV文件路径

    Returns:
        pd.DataFrame: 读取的数据框
    """
    data = pd.read_csv(path, header=0, sep=",", encoding="utf-8")
    return data


def outputdata(path, data, is_index=False):
    """
    输出数据到CSV文件

    Args:
        path (str): 输出文件路径
        data (pd.DataFrame): 要输出的数据框
        is_index (bool): 是否包含索引列
    """
    data.to_csv(path, index=is_index, header=True, sep=",", mode="w", encoding="utf-8")


def transcolname(df, column_mapping):
    """
    转换列名为英文

    Args:
        df (pd.DataFrame): 输入数据框
        column_mapping (dict): 列名映射字典

    Returns:
        pd.DataFrame: 列名转换后的数据框
    """
    df.rename(columns=column_mapping, inplace=True)
    return df


def trans_datetime(df):
    """
    日期时间特征转换（与训练时保持一致）

    Args:
        df (pd.DataFrame): 包含Date列的数据框

    Returns:
        pd.DataFrame: 添加了时间特征的数据框
    """
    ret_df = pd.DataFrame()
    dt = df["Date"]

    # 提取年月日特征
    ret_df["year"] = dt.transform(lambda x: int(x.split("-")[0]))
    ret_df["month"] = dt.transform(lambda x: int(x.split("-")[1]))
    ret_df["day"] = dt.transform(lambda x: int(x.split("-")[2][:2]))

    # 合并时间特征
    df = pd.concat([df, ret_df], axis=1)

    # 将日期转换为序号
    unique_dates = pd.Series(df["Date"].unique()).sort_values().reset_index(drop=True)
    date_mapping = {date: rank + 1 for rank, date in enumerate(unique_dates)}
    df["Date"] = df["Date"].map(date_mapping)

    return df


def processing_feature_test():
    """
    测试数据特征处理函数
    读取测试数据并进行与训练数据相同的预处理

    Returns:
        pd.DataFrame: 处理后的测试数据
    """
    # 读取测试数据
    data = inputdata(TEST_DATA_PATH)

    # 列名映射（与训练时保持一致）
    column_mapping = {
        "股票代码": "StockCode",
        "日期": "Date",
        "开盘": "Open",
        "收盘": "Close",
        "最高": "High",
        "最低": "Low",
        "成交量": "Volume",
        "成交额": "Turnover",
        "振幅": "Amplitude",
        "涨跌额": "PriceChange",
        "换手率": "TurnoverRate",
        "涨跌幅": "PriceChangePercentage",
    }

    # 转换列名
    data = transcolname(data, column_mapping)

    # 删除涨跌幅列
    data.drop(columns=["PriceChangePercentage"], inplace=True)

    # 日期特征转换
    data = trans_datetime(data)

    # 只保留最近pred_len天的数据用于预测
    max_date = data["Date"].max()
    data = data[data["Date"] > max_date - pred_len]

    return data


# 处理测试数据
data = processing_feature_test()
max_date = data["Date"].max()

# 导入PyTorch相关库
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader

# 模型参数（与训练时保持一致）
lstm_config = MODEL_PARAMS["LSTM"]
training_config = TRAINING_PARAMS

input_size = 1
hidden_size = lstm_config["hidden_size"]
num_layers = lstm_config["num_layers"]
output_size = 1
dropout = lstm_config["dropout"]

# 设备配置
if training_config["device"] == "auto":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
else:
    device = torch.device(training_config["device"])


# LSTM模型定义（与训练时保持一致）
class LSTMModel(nn.Module):
    """
    LSTM时间序列预测模型
    """

    def __init__(self, input_size, hidden_size, num_layers, output_size, dropout):
        super(LSTMModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(
            input_size, hidden_size, num_layers, batch_first=True, dropout=dropout
        )
        self.fc = nn.Linear(hidden_size, output_size)

    def __del__(self):
        del self.hidden_cell

    def forward(self, x):
        out, self.hidden_cell = self.lstm(x, self.hidden_cell)
        out = self.fc(out[:, -1, :])
        return out


# 定义需要预测的列
PREDICT_COLS = [
    "Open",           # 开盘价
    "Close",          # 收盘价
    "High",           # 最高价
    "Low",            # 最低价
    "Volume",         # 成交量
    "Turnover",       # 成交额
    "Amplitude",      # 振幅
    "PriceChange",    # 涨跌额
    "TurnoverRate",   # 换手率
]

# 获取列名信息
column_names = data.columns.tolist()
colname2index = {x: i for i, x in enumerate(column_names)}
stockcodes = data["StockCode"].drop_duplicates().tolist()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
all_preds = []  # 存储所有股票的预测结果


def process_data_predict(data, stockcode):
    """
    为单支股票准备预测数据

    Args:
        data (pd.DataFrame): 测试数据
        stockcode (str): 股票代码

    Returns:
        torch.Tensor: 准备好的预测数据张量
    """
    # 筛选单支股票的数据
    data = data[data["StockCode"] == stockcode]
    data = np.array([data.values])
    return torch.tensor(data, dtype=torch.float32).to(device)


# 对每支股票进行预测
for stockcode in stockcodes:
    temp_preds = {}

    # 准备当前股票的预测数据
    predict_data = process_data_predict(data, stockcode).to(device)

    # 加载训练好的Close价格预测模型
    model_name = f"{MODEL_DIR}/model_Close.bin"
    estimator = pickle.load(open(model_name, "rb"))

    # 初始化模型的隐藏状态
    estimator.hidden_cell = (
        torch.zeros(
            estimator.num_layers, predict_data.size(0), estimator.hidden_size
        ).to(device),
        torch.zeros(
            estimator.num_layers, predict_data.size(0), estimator.hidden_size
        ).to(device),
    )

    # 进行预测（不计算梯度）
    with torch.no_grad():
        predi = estimator(predict_data)

    # 获取预测结果
    pred = predi[-1].cpu().detach().numpy()

    # 存储预测结果
    all_preds.append((stockcode, pred))


# 计算涨跌幅并排序
pricechangerate = []
for i in range(len(all_preds)):
    stockcode, pred = all_preds[i]

    # 获取该股票当前的收盘价
    preClose = data[(data["StockCode"] == stockcode) & (data["Date"] == max_date)][
        "Close"
    ].values[0]

    # 计算预测涨跌幅 = (预测价格 - 当前价格) / 当前价格 * 100%
    pricechangerate.append((stockcode, (pred - preClose) / preClose * 100))

# 按涨跌幅排序（降序）
pricechangerate = sorted(pricechangerate, key=lambda x: x[1], reverse=True)

# 获取涨幅最大的前10支股票和涨幅最小的后10支股票
pred_top_10_max_target = [x[0] for x in pricechangerate[:10]]   # 涨幅最大的10支
pred_top_10_min_target = [x[0] for x in pricechangerate[-10:]]  # 涨幅最小的10支

# 构建结果数据
data = {
    "涨幅最大股票代码": pred_top_10_max_target,
    "涨幅最小股票代码": pred_top_10_min_target,
}

# 输出预测结果到CSV文件
df = pd.DataFrame(data)
df.to_csv(RESULT_PATH, index=False)
